from PIL import Image, ImageFile

from torch.utils.data import Dataset # 用来构建自定义数据集的基类
import os.path as osp # 方便进行文件路径相关的操作
import random
import torch
ImageFile.LOAD_TRUNCATED_IMAGES = True # 当图像文件损坏或未完全下载时，它可能是一个“截断的”图像。默认情况下，Pillow 在尝试打开这样的文件时会抛出 IOError。将这个标志设置为 True，是告诉 Pillow 尝试加载这些不完整的图像，而不是直接报错。

# 以一种健壮的方式读取图像文件。它被设计用来处理在大量输入/输出（I/O）操作中可能发生的 IOError，并通过持续重试来确保图像最终被成功读取。
def read_image(img_path):
    """Keep reading image until succeed.
    This can avoid IOError incurred by heavy IO process."""
    got_img = False
    if not osp.exists(img_path): # 检查文件路径是否存在
        raise IOError("{} does not exist".format(img_path)) # 路径不存在就抛出异常
    while not got_img: # 只要没有收到图片，就持续循环
        try:
            img = Image.open(img_path).convert('RGB') # 将图像转换为 RGB 颜色模式
            got_img = True
        except IOError:
            print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
            pass
    return img


class BaseDataset(object):
    """
    Base class of reid dataset
    """

    def get_imagedata_info(self, data):
        pids, cams, tracks = [], [], [] # 分别用来存放所有的行人ID、摄像头ID和轨迹ID。

        for _, pid, camid, trackid in data:
            pids += [pid]
            cams += [camid]
            tracks += [trackid]
        pids = set(pids)
        cams = set(cams)
        tracks = set(tracks)
        num_pids = len(pids)
        num_cams = len(cams)
        num_imgs = len(data)
        num_views = len(tracks)
        return num_pids, num_imgs, num_cams, num_views

    def print_dataset_statistics(self): # 这个方法在基类中只是一个“框架”或“接口”，本身没有具体的实现逻辑，任何继承 BaseDataset 的子类都 必须 自己去实现（重写）这个方法，否则在调用它时就会抛出 NotImplementedError 异常
        raise NotImplementedError


class BaseImageDataset(BaseDataset):
    """
    Base class of image reid dataset
    """

    def print_dataset_statistics(self, train, query, gallery):
        num_train_pids, num_train_imgs, num_train_cams, num_train_views = self.get_imagedata_info(train)
        num_query_pids, num_query_imgs, num_query_cams, num_query_views = self.get_imagedata_info(query)
        num_gallery_pids, num_gallery_imgs, num_gallery_cams, num_gallery_views = self.get_imagedata_info(gallery)

        print("Dataset statistics:")
        print("  ----------------------------------------")
        print("  subset   | # pids | # images | # cameras")
        print("  ----------------------------------------")
        print("  train    | {:5d} | {:8d} | {:9d}".format(num_train_pids, num_train_imgs, num_train_cams))
        print("  query    | {:5d} | {:8d} | {:9d}".format(num_query_pids, num_query_imgs, num_query_cams))
        print("  gallery  | {:5d} | {:8d} | {:9d}".format(num_gallery_pids, num_gallery_imgs, num_gallery_cams))
        print("  ----------------------------------------")

# 这个类继承自 torch.utils.data.Dataset。这是 PyTorch 框架中用于构建自定义数据集的标准方式。任何继承自 Dataset 的类，都必须实现 __len__ 和 __getitem__ 这两个魔法方法。
# 这个类的实例可以被 PyTorch 的 DataLoader 直接使用，从而实现高效的数据加载、批处理（batching）、打乱（shuffling）和多线程加载。
class ImageDataset(Dataset):
    def __init__(self, dataset, transform=None):
        self.dataset = dataset # 接收一个包含所有数据信息的列表（比如 train 列表，其中每个元素是 (img_path, pid, camid, trackid)）
        self.transform = transform #  接收一个可选的 transform 参数。在深度学习中，transform 通常是一个函数或一系列操作（例如，来自 torchvision.transforms），用于对图像进行预处理和数据增强（如缩放、裁剪、翻转、归一化等）。

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        img_path, pid, camid, trackid = self.dataset[index]
        img = read_image(img_path)

        if self.transform is not None:
            img = self.transform(img)

        return img, pid, camid, trackid,img_path.split('/')[-1] # 返回图像文件名